import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.mplot3d import Axes3D
import os
import numpy as np
import itertools
from matplotlib.colors import TwoSlopeNorm
# import matplotlib.font_manager as font_manager
# from scipy.interpolate import interp1d


class plots(object):

    def __init__(self, **kwargs):
        self.save = kwargs.get('save', False)
        self.wd = kwargs.get('wd', None)
        self.fileName = kwargs.get('fileName', None)
        self.dpi = kwargs.get('dpi', 800)
        self.fmt = kwargs.get('fmt', 'png')
        # SIZING
        self.figsize = kwargs.get('figsize', [5.76, 4.32])
        self.fontsize = kwargs.get('fontsize',  14)
        self.legendFontsize = kwargs.get('legendFontsize', 14)
        self.axisFontSize = kwargs.get('axisFontSize', 14)
        self.xtickFontsize = kwargs.get('xtickFontsize', 2 * (self.axisFontSize/2 - 1))
        self.ytickFontsize = kwargs.get('ytickFontsize', 2 * (self.axisFontSize/2 - 1))
        # self.lableFontWeight = kwargs.get('lableFontWeight', 'normal')
        # self.legendFontWeight = kwargs.get('legendFontWeight', 'bold')

        # SETTING RCPARAMS
        plt.rcParams['figure.figsize'] = self.figsize
        plt.rcParams.update({'font.size': self.fontsize})
        plt.rcParams['legend.fontsize'] = self.legendFontsize
        plt.rcParams['xtick.labelsize'] = self.xtickFontsize
        plt.rcParams['ytick.labelsize'] = self.ytickFontsize
        plt.rcParams['font.family'] = 'sans-serif'
        # plt.rcParams['font.sans-serif'] = 'Whitney'
        # plt.rcParams['font.weight'] = 350
        plt.rcParams['font.sans-serif'] = 'Merriweather Sans'
        plt.rcParams['lines.markeredgecolor'] = '#000000'
        plt.rcParams['axes.unicode_minus'] = False

    def imagePlot(self, **kwargs):

        cmap = kwargs.get('cmap', 'viridis')
        interpolation = kwargs.get('interpolation', None)
        vmax = kwargs.get('vmax', None)
        vmin = kwargs.get('vmin', None)
        x = kwargs.get('x', np.ones((10, 10)))
        cbartitle = kwargs.get('cbartitle', None)
        title = kwargs.get('title', None)

        fig, ax = plt.subplots()
        ax.set_aspect('equal')

        im = ax.imshow(x, extent=[-x.shape[1] / 2., x.shape[1] / 2., -x.shape[0] / 2., x.shape[0] / 2.],
                       interpolation=interpolation,
                       cmap=cmap, vmin=vmin, vmax=vmax)

        ax.minorticks_on()
        ax.tick_params(which='minor', direction='in', length=4)
        ax.tick_params(which='major', direction='in', length=8)

        # divider = make_axes_locatable(ax)
        # cax = divider.append_axes("right", size="5%", pad=0.2)
        # cbar = plt.colorbar(im, cax=cax)
        # cbar.ax.minorticks_on()
        # cbar.ax.tick_params(which='minor', direction='in', length=2)
        # cbar.ax.tick_params(direction='in', length=4)
        #
        # if title:
        #     ax.set_title()
        #
        # if cbartitle:
        #     cbar.ax.get_yaxis().labelpad = 20
        #     cbar.ax.set_ylabel('{}'.format(cbartitle), rotation=270)

        if self.save:
            plt.savefig(os.path.join(self.wd, '{}.{}'.format(self.fileName, self.fmt)), dpi=self.dpi)
            plt.close()
        else:
            plt.show()

        plt.close()

    def categoricalScatter(self, **kwargs):
        xnew = kwargs.get('x', np.ones((3, 3)))
        dim = kwargs.get('dim', 3)
        labels = kwargs.get('labels', None)
        cdict = kwargs.get('cdict', None)
        labelpad = kwargs.get('labelpad', 0)
        labelsize = kwargs.get('labelsize', None)
        xlim = kwargs.get('xlim', (None, None))
        ylim = kwargs.get('ylim', (None, None))
        zlim = kwargs.get('zlim', (None, None))
        xlabel = kwargs.get('xlabel', 'PC1')
        ylabel = kwargs.get('ylabel', 'PC2')
        zlabel = kwargs.get('zlabel', 'PC3')
        title = kwargs.get('title', '')
        legendorder = kwargs.get('lengendorder', None)

        fig = plt.figure()
        fig.tight_layout()

        if dim == 3:
            ax = Axes3D(fig)
            for label in np.unique(labels):
                i = np.where(labels == label)
                ax.scatter(xnew[i, 0], xnew[i, 1], xnew[i, 2], c=cdict[label], s=60, label=label, depthshade=False)

            if legendorder:
                handles, labels = ax.get_legend_handles_labels()
                _handles = []
                _labels = []
                for i in legendorder:
                    _handles = _handles + [handles[i]]
                    _labels = _labels + [labels[i]]
                ax.legend(_handles, _labels, loc='best')
            else:
                ax.legend(loc='best')

            ax.set_xlabel("{}".format(xlabel))
            ax.set_ylabel("{}".format(ylabel))
            ax.set_zlabel("{}".format(zlabel))

            ax.set_xlim(xlim)
            ax.set_ylim(ylim)
            ax.set_zlim(zlim)

            ax.xaxis.labelpad = labelpad
            ax.yaxis.labelpad = labelpad
            ax.zaxis.labelpad = labelpad

            ax.tick_params(which='major', direction='in', length=8, labelsize=labelsize)
            ax.set_title("{}".format(title))

        elif dim == 2:
            ax = plt.subplot()
            for label in np.unique(labels):
                i = np.where(labels == label)
                ax.scatter(xnew[i, 0], xnew[i, 1], c=cdict[label], s=60, label=label)

            if legendorder:
                handles, labels = ax.get_legend_handles_labels()
                _handles = []
                _labels = []
                for i in legendorder:
                    _handles = _handles + [handles[i]]
                    _labels = _labels + [labels[i]]
                ax.legend(_handles, _labels, loc='best')
            else:
                ax.legend(loc='best')

            ax.set_xlabel("{}".format(xlabel))
            ax.set_ylabel("{}".format(ylabel))

            ax.set_xlim(xlim)
            ax.set_ylim(ylim)

            ax.xaxis.labelpad = labelpad
            ax.yaxis.labelpad = labelpad

            ax.minorticks_on()
            ax.tick_params(which='minor', direction='in', length=4)
            ax.tick_params(which='major', direction='in', length=8, labelsize=labelsize)
            ax.set_title("{}".format(title))

        else:
            print('输入错误')

        if self.save:
            plt.savefig(os.path.join(self.wd, '{}.{}'.format(self.fileName, self.fmt)), dpi=self.dpi)
            plt.close()
        else:
            plt.show()

    def linePlot(self, **kwargs):
        x = kwargs.get('x', 0)
        y = kwargs.get('y', 0)
        xlim = kwargs.get('xlim', (None, None))
        ylim = kwargs.get('ylim', (None, None))
        xlabel = kwargs.get('xlabel', '')
        ylabel = kwargs.get('ylabel', '')
        xticks = kwargs.get('xticks', None)
        cmap = kwargs.get('cmap', plt.cm.Set1)
        label = kwargs.get('label', None)
        markersize = kwargs.get('markersize', 0)
        marker = kwargs.get('marker', itertools.cycle(
            ('o', 'v', '^', '<', '>', 's', '8', 'p', '*', 'h', 'd', 'H', 'D', 'P', 'X')))
        # linestyle = kwargs.get('linestyle', itertools.cycle(('-', '--', '-.', ':')))
        linewidth = kwargs.get('linewidth', 1)
        minorticks = kwargs.get('minorticks', True)
        gridOn = kwargs.get('gridOn', True)

        fig, ax = plt.subplots()

        if isinstance(y[0], np.ndarray):
            for i in range(y.shape[1]):
                if isinstance(x[0], np.ndarray):
                    ax.plot(x[:, i], y[:, i], color=cmap(i), label=label[i], marker=marker.__next__(),
                            markersize=markersize, linewidth=linewidth)
                else:
                    ax.plot(x, y[:, i], color=cmap(i), label=label[i], marker=marker.__next__(), markersize=markersize,
                            linewidth=linewidth)
        else:
            if marker == itertools.cycle(('o', 'v', '^', '<', '>', 's', '8', 'p')):
                marker = None
            else:
                pass
            ax.plot(x, y, marker=marker, label=label, linewidth=linewidth)

        # SETTING THE LEGEND
        ax.legend(loc='best')
        # ax.legend(prop=font_manager.FontProperties(weight=self.legendFontWeight))

        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.set_xlabel('{}'.format(xlabel))
        ax.set_ylabel('{}'.format(ylabel))

        if minorticks:
            ax.minorticks_on()
            ax.tick_params(which='minor', direction='in', length=4)
        ax.tick_params(which='major', direction='in', length=8)

        if xticks.any():
            plt.xticks(xticks)
        if gridOn:
            ax.grid(linestyle='--')

        fig.tight_layout()

        if self.save:
            plt.savefig(os.path.join(self.wd, '{}.{}'.format(self.fileName, self.fmt)), dpi=self.dpi)
        else:
            plt.show()

    def arcldPlot(self, **kwargs):
        xbin = kwargs.get('xbin', None)
        z = kwargs.get('z', None)
        cmap = kwargs.get('cmap', 'viridis')
        vmax = kwargs.get('vmax', None)
        vmin = kwargs.get('vmin', None)
        centered = kwargs.get('centered', False)
        levels = kwargs.get('levels', 128)

        [r, t] = np.meshgrid(xbin, np.linspace(0, 2 * np.pi, num=360))
        x = r * np.cos(t)
        y = r * np.sin(t)

        arcldArray = np.vstack((z, z))
        fig, ax = plt.subplots()
        if not centered:
            im = ax.contourf(x, y, arcldArray, levels, vmin=vmin, vmax=vmax, cmap=cmap)
        else:
            im = ax.contourf(x, y, arcldArray, levels, vmin=vmin, vmax=vmax, cmap=cmap, norm=TwoSlopeNorm(0))

        ax.set_aspect('equal')
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.2)
        cbar = plt.colorbar(im, cax=cax)
        cbar.ax.tick_params(which='minor', direction='in', length=2)
        cbar.ax.tick_params(direction='in', length=4)
        cbar.ax.set_ylabel('Relative Frequency', rotation=270)
        cbar.ax.get_yaxis().labelpad = 20

        ax.set_xlabel("Chord Length (μm)")
        ax.set_ylabel("Chord Length (μm)")
        # im.set_clim(vmin=0)

        if self.save:
            plt.savefig(os.path.join(self.wd, '{}.{}'.format(self.fileName, self.fmt)), dpi=self.dpi)
            plt.close()
        else:
            plt.show()

    def histPlot(self, **kwargs):
        x = kwargs.get('x', None)
        binWidth = kwargs.get('binWidth', False)
        numBins = kwargs.get('numBins', False)
        normed = kwargs.get('normed', 'density')
        xlim = kwargs.get('xlim', (None, None))
        ylim = kwargs.get('ylim', (None, None))
        xlabel = kwargs.get('xlabel', '')
        ylabel = kwargs.get('ylabel', '')
        label = kwargs.get('label', None)
        cmap = kwargs.get('cmap', plt.cm.Set1)
        linewidth = kwargs.get('linewidth', 1)
        edgecolor = kwargs.get('edgecolor', False)

        _dim = len(x)
        fig, ax = plt.subplots()
        patches = []
        if isinstance(x[0], np.ndarray):
            for i in range(_dim):
                if binWidth:
                    _binWidth = binWidth
                    _numBins = int((max(x[i]) - min(x[i])) // _binWidth)
                else:
                    _numBins = numBins
                _bins = np.linspace(min(x[i]), max(x[i]), _numBins)

                if normed == 'density':
                    _n, _b, _p = ax.hist(x[i], bins=_bins, density=True, label=label[i])
                    patches.append(_p)
                elif normed == 'probability':
                    _weights = np.ones_like(x[i]) / len(x[i])
                    _n, _b, _p = plt.hist(x[i], bins=_bins, weights=_weights, label=label[i])
                    patches.append(_p)

            # CREATE COLOR LIST

        else:
            if binWidth:
                _binWidth = binWidth
                _numBins = int((max(x) - min(x)) // _binWidth)
            else:
                _numBins = numBins
            _bins = np.linspace(min(x), max(x), _numBins)

            if normed == 'density':
                _n, _b, _p = ax.hist(x, bins=_bins, density=True, label=label)
                patches.append(_p)
            elif normed == 'probability':
                _weights = np.ones_like(x) / len(x)
                _n, _b, _p = plt.hist(x, bins=_bins, weights=_weights, label=label)
                patches.append(_p)

        # PATCH PROPERTIES
        if not edgecolor:
            edgecolor = ['w'] * _dim

        for i in range(_dim):
            for _bar in patches[i]:
                _bar.set_facecolor(cmap[i])
                _bar.set_edgecolor(edgecolor[i])

        # AXES PROPERTIES
        ax.set_xlim(xlim)
        ax.set_ylim(ylim)
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.minorticks_on()
        ax.tick_params(which='minor', direction='in', length=3)
        ax.tick_params(which='major', direction='in', length=6)
        ax.set_axisbelow(True)
        ax.grid(linestyle='--', linewidth=0.5, fillstyle='bottom')

        # this method to assign legend label is discarded
        # ax.legend([_p for _p in patches], label)
        ax.legend(loc='upper center',
                  facecolor='w',
                  edgecolor='#D6D6D6',
                  framealpha=1.0)

        plt.tight_layout()

        plt.show()

    def violinPlot(self):
        pass

    def contourPlot(self):
        pass


